
#include "fun.h"
#include "hs_cuda_scan.cuh"
#include "hs_cuda_elemwise.cuh"
#include "hs_cuda_reduce.cuh"
#include "hs_cuda_inputwise.cuh"

void getPrefixOpenSum(const hs::raster::RasterData<float> & raster, hs::raster::RasterData<float> & prefix_sum)
{
    /* 将数据拷贝到GPU端 */
    hscuda::raster::RasterData<float> device_raster;
    hscuda::raster::host2device(raster.view(), device_raster);

    int N = raster.view().numel();
    int size_x = raster.view().sizeX();
    int size_y = raster.view().sizeY();

    /* 为输出数据分配显存 */
    hscuda::raster::RasterData<float> device_prefix_sum(size_x, size_y);

    /* 直接开扫描调用模板函数即可 */
    hs::cuda::scan::PrefixOpenSum(N
        , device_raster.view().pData()
        , device_prefix_sum.view().pData());
    cudaDeviceSynchronize();

    /* 将结果拷贝到主机端 */
    device_prefix_sum.toHost(prefix_sum);
}

/* 应用 elemwise, 提取大于thre的数据掩膜 */
struct GenMaskFunctor
{
    GenMaskFunctor(float thre) : m_thre(thre) { }

    __device__ float operator()(float val) const
    {
        if (val > m_thre) {
            return 1.0f;
        }
        else {
            return 0.0f;
        }
    }

private:
    float m_thre;
};

/* 应用 reduction, 计算掩膜的和, 此值为输出结果的长度 */
struct MaskSumFunctor
{
    __device__ __forceinline__ float operator()(float x, float y) const
    {
        return x + y;
    }
};

/* 应用 inputwise, 遍历输入掩膜和前缀和, 掩膜元素为1是目标元素, 其在输出位置索引通过前缀和读取 */
struct MaskIndexFunctor
{
    MaskIndexFunctor(int out_size, float* p_out_datas)
        : m_out_size(out_size), m_p_out_datas(p_out_datas)
    { }

    __device__ void operator()(float mask, float fid, float value)
    {
        int id = (int)floor(fid);
        if (id < m_out_size && id >= 0 && mask != 0) {
            m_p_out_datas[id] = value;
        }
    }

private:
    int m_out_size;
    float* m_p_out_datas;
};

void extractElements(const hs::raster::RasterData<float> & raster, float thre, hs::raster::RasterData<float> & elems)
{
    /* 将主机数据复制到GPU */
    hscuda::raster::RasterData<float> device_raster;
    hscuda::raster::host2device(raster.view(), device_raster);

    /* 应用 elemwise, 提取大于thre的数据掩膜 */
    int N = device_raster.view().numel();
    hscuda::raster::RasterData<float> mask(device_raster.view().sizeX(), device_raster.view().sizeY());
    hs::cuda::elemwise::Unary(N
        , device_raster.view().pData()
        , GenMaskFunctor(thre)
        , mask.view().pData()
    );
    cudaDeviceSynchronize();

    {
        printf("Mask:\n");
        hs::raster::RasterData<float> host_mask;
        mask.toHost(host_mask);
        printmat(host_mask.view());
    }

    /* 应用 reduction, 计算掩膜的和, 此值为输出结果的长度 */
    float res_sum = 0.0f;
    hs::cuda::reduce::apply(N
        , mask.view().pData()
        , MaskSumFunctor() 
        , res_sum          
        , &res_sum);
    cudaDeviceSynchronize();
    printf("total number of elements: %d\n", (int)res_sum);


    /* 应用 scan, 求解掩膜的开扫描前缀和, 开扫描前缀和即是输出元素的索引 */    
    hscuda::raster::RasterData<float> prefix_sum(mask.view().sizeX(), mask.view().sizeY());
    hs::cuda::scan::PrefixOpenSum(N
        , mask.view().pData()
        , prefix_sum.view().pData());
    cudaDeviceSynchronize();
    {
        hs::raster::RasterData<float> host_prefix_sum;
        prefix_sum.toHost(host_prefix_sum);
        printf("Prefix sum:\n");
        printmat(host_prefix_sum.view());
    }

    /* 应用 inputwise, 遍历输入掩膜和前缀和, 获取掩膜元素为1以及在同样索引位置的前缀和(该前缀和恰好为该元素在输出数组的索引) */
    // 输出数据分配显存
    hscuda::raster::RasterData<float> device_elems((int)res_sum, 1);

    hs::cuda::inputwise::Ternary(N
        , mask.view().pData()
        , prefix_sum.view().pData()
        , device_raster.view().pData()
        , MaskIndexFunctor(device_elems.view().numel()
                         , device_elems.view().pData())
    );
    cudaDeviceSynchronize();
    
    /* 将结果拷贝到CPU */
    device_elems.toHost(elems);
}
